import svox2
import svox2.utils
import argparse
import torch
import cv2
from tqdm import tqdm
import os



def main(args, grid):
        size = grid.shape[0]
        if args.axis == 'x':
            x = torch.ones(size, size, 1, device=args.device) * args.axis_index
            y = torch.arange(size, device=args.device).repeat(size, 1).unsqueeze(-1)
            z = torch.arange(size, device=args.device).repeat(size, 1).transpose(0, 1).unsqueeze(-1)
            xyz = torch.stack([x, y, z], dim=-1)
        elif args.axis == 'y':
            x = torch.arange(size, device=args.device).repeat(size, 1).unsqueeze(-1)
            y = torch.ones(size, size, 1, device=args.device) * args.axis_index
            z = torch.arange(size, device=args.device).repeat(size, 1).transpose(0, 1).unsqueeze(-1)
            xyz = torch.stack([x, y, z], dim=-1)
        elif args.axis == 'z':
            x = torch.arange(size, device=args.device).repeat(size, 1).unsqueeze(-1)
            y = torch.arange(size, device=args.device).repeat(size, 1).transpose(0, 1).unsqueeze(-1)
            z = torch.ones(size, size, 1, device=args.device) * args.axis_index
            xyz = torch.stack([x, y, z], dim=-1)

        xyz = xyz.reshape(-1, 3)

        # forward
        xyz = xyz.to(args.device)

        with torch.no_grad():
            density, rgb = grid.sample(xyz, use_kernel=False, grid_coords=True)

        # to image
        density = density.reshape(size, size)
        sigmoid_density = torch.sigmoid(density)
        sigmoid_density = sigmoid_density.detach().cpu().numpy()
        sigmoid_density = (sigmoid_density * 255).astype('uint8')

        # save
        cv2.imwrite(os.path.join(args.out, 'density_{}_{}.png'.format(args.axis, args.axis_index)), sigmoid_density)
    
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    # device
    parser.add_argument('--device', type=str, default='cpu')
    # path
    parser.add_argument('--ckpt', type=str, default='outputs/blender_lego_radius/ckpt.npz')
    # axis
    parser.add_argument('--axis', type=str, default='z')
    # axis index
    parser.add_argument('--axis_index', type=int, default=256)
    # output path
    parser.add_argument('--out', type=str, default='outputs/plane_visualize')
    # end
    args = parser.parse_args()

    os.makedirs(args.out, exist_ok=True)

    grid = svox2.SparseGrid.load(args.ckpt, device=args.device)
    grid._scaling = 1.0 
    grid._offset = 0.0

    for i in tqdm(range(100, 400)):
        args.axis_index = i
        main(args, grid)